#include "SectorMon.h"

BOOLEAN StartFilter = FALSE;
BOOLEAN InitFilter = FALSE;
NTFS_BOOT_SEC BootSec;
GLOBAL_CONTEXT GlobalContext = {0};
USHORT BytesPerSector = 512;


NTSTATUS
RecoverBootStatus()
{
	NTSTATUS ns;
	HANDLE FileHandle;
	OBJECT_ATTRIBUTES oa;
	IO_STATUS_BLOCK Iob;
	LARGE_INTEGER Offset = {0};
	UCHAR Buffer[16];
	DECLARE_CONST_UNICODE_STRING(FileName, L"\\??\\C:\\Windows\\bootstat.dat");

	InitializeObjectAttributes(&oa, (PUNICODE_STRING)&FileName, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, NULL);

	ns = ZwCreateFile(&FileHandle, 
		FILE_READ_DATA | FILE_WRITE_DATA | SYNCHRONIZE, 
		&oa, 
		&Iob,
		NULL,
		0,
		FILE_SHARE_READ,
		FILE_OPEN, 
		FILE_SYNCHRONOUS_IO_NONALERT,
		NULL, 
		0);

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	ns = ZwReadFile(FileHandle, 
		NULL, 
		NULL, 
		NULL, 
		&Iob, 
		Buffer, 
		16, 
		&Offset, 
		NULL);

	if (!NT_SUCCESS(ns)) {

		ZwClose(FileHandle);
		return ns;
	}

	if (Buffer[10] == 0) {

		Buffer[10] = 1;
	}

	ns = ZwWriteFile(FileHandle, 
		NULL, 
		NULL, 
		NULL, 
		&Iob, 
		Buffer, 
		16, 
		&Offset, 
		NULL);

	ZwClose(FileHandle);
	return ns;
}


BOOLEAN
IsRestoreDisk()
{
	NTSTATUS ns;
	UNICODE_STRING SecMonName;
	UNICODE_STRING RestoreDiskName;
	OBJECT_ATTRIBUTES oa;
	HANDLE SecMonHandle;
	UCHAR ValueBuffer[128];
	ULONG RetLength;
	PKEY_VALUE_PARTIAL_INFORMATION Value = (PKEY_VALUE_PARTIAL_INFORMATION)ValueBuffer;


	RtlInitUnicodeString(&SecMonName, L"\\Registry\\Machine\\System\\CurrentControlSet\\Control\\SecMon");
	InitializeObjectAttributes(&oa,
		&SecMonName,
		OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
		NULL,
		NULL);
	ns = ZwOpenKey(&SecMonHandle,
		KEY_READ,
		&oa);
	if (!NT_SUCCESS(ns)) {

		return FALSE;
	}

	RtlInitUnicodeString(&RestoreDiskName, L"Current");
	ns = ZwQueryValueKey(SecMonHandle,
		&RestoreDiskName,
		KeyValuePartialInformation,
		ValueBuffer,
		sizeof(ValueBuffer),
		&RetLength);

	ZwClose(SecMonHandle);

	if (!NT_SUCCESS(ns)) {

		return FALSE;
	}
	
	if (*((ULONG *)Value->Data) == 0) {

		return FALSE;
	}
	
	return TRUE;
}


NTSTATUS
RestoreDisk(
			__in PDISK_FLT_DEVICE_EXT DevExt
			)
{
	ULONGLONG BitCount;
	PVOID BitBuf = GlobalContext.MapBuffer;
	ULONG BitSize = GlobalContext.MapSize;
	LARGE_INTEGER Offset;
	ULONGLONG i;
	LONG Ret;
	PVOID SectorBuf;
	NTSTATUS ns = STATUS_SUCCESS;
	IO_STATUS_BLOCK Iob;
	HANDLE FileHandle = GlobalContext.SparseFileHandle;
	PDEVICE_OBJECT PhyDevObj = DevExt->PhysicalDeviceObject;
	FILE_DISPOSITION_INFORMATION FileDisInfo;
	LARGE_INTEGER Mm3Second = {(ULONG)(-3 * 1000 * 1000 * 10), -1};

	BitCount = BitSize << 3;
	SectorBuf = ExAllocatePoolWithTag(PagedPool, BytesPerSector, 'rsdk');


	for (i = 0; i < BitCount; i++) {

		if (i % 0x1000 == 0) {

			KdPrint(("Per : %I64d / %I64d\n", i, BitCount));
		}

		Ret = GetBit(BitBuf, BitSize, i);

		ASSERT(Ret != -1);
		if (Ret == -1) {

			ns = STATUS_UNSUCCESSFUL;
			break;
		}
		else if (Ret == 1) {

			Offset.QuadPart = i * BytesPerSector;
			ns = ZwReadFile(FileHandle, 
				NULL, 
				NULL, 
				NULL, 
				&Iob, 
				SectorBuf, 
				BytesPerSector, 
				&Offset, 
				NULL);

			ASSERT(NT_SUCCESS(ns));

			if (!NT_SUCCESS(ns)) {

				break;
			}

			ns = DirectReadWrite(PhyDevObj, 
				DIRECT_WRITE, 
				SectorBuf, 
				&Offset, 
				BytesPerSector);

			ASSERT(NT_SUCCESS(ns));

			if (!NT_SUCCESS(ns)) {

				break;
			}
		}
	}

	ExFreePoolWithTag(SectorBuf, 'rsdk');

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	Offset.QuadPart = 0;
	RtlZeroMemory(BitBuf, BitSize);
	ns = ZwWriteFile(GlobalContext.FileHandle, 
		NULL, 
		NULL, 
		NULL, 
		&Iob, 
		BitBuf, 
		BitSize, 
		&Offset, 
		NULL);

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	FileDisInfo.DeleteFile = TRUE;
	ns = ZwSetInformationFile(FileHandle, 
		&Iob, 
		&FileDisInfo, 
		sizeof(FILE_DISPOSITION_INFORMATION), 
		FileDispositionInformation);

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	ZwClose(FileHandle);

	ns = RecoverBootStatus();
	
	KeDelayExecutionThread(KernelMode, FALSE, &Mm3Second);

	KeBugCheck(POWER_FAILURE_SIMULATE);
}


VOID 
WriteWorkThread( 
				__in PVOID  Context
				)
{
	PWRITE_ELEMENT ReqEntry = NULL;
	PDISK_FLT_DEVICE_EXT DevExt = (PDISK_FLT_DEVICE_EXT)Context;
	IO_STATUS_BLOCK Iob;

	KeSetPriorityThread(KeGetCurrentThread(), LOW_REALTIME_PRIORITY);

	for(;;) {	

		KeWaitForSingleObject(
			&DevExt->ReqEvent,
			Executive,
			KernelMode,
			FALSE,
			NULL
			);

		if (DevExt->ThreadFlag) {
			PsTerminateSystemThread(STATUS_SUCCESS);
			return;
		}

		while ((ReqEntry = (PWRITE_ELEMENT)ExInterlockedRemoveHeadList(&DevExt->ReqList,
					&DevExt->ReqLock)) != NULL) {

				ZwWriteFile(ReqEntry->File, 
					NULL, 
					NULL, 
					NULL, 
					&Iob, 
					ReqEntry->Buffer, 
					ReqEntry->Length, 
					&ReqEntry->Offset, 
					NULL);
				if (ReqEntry->NeedFree) {

					ExFreeToPagedLookasideList(&DevExt->SectorBufHeader, ReqEntry->Buffer);
				}

				ExFreeToPagedLookasideList(&DevExt->WriteElemHeader, ReqEntry);
		}
	}
}


NTSTATUS 
SkipAndCallDriver(
				  __in PDEVICE_OBJECT DeviceObject,
				  __in PIRP Irp 
				  )
{
	IoSkipCurrentIrpStackLocation(Irp);
	return IoCallDriver(DeviceObject, Irp);
}


NTSTATUS
SyncIrpCompletionRoutine(
						 __in PDEVICE_OBJECT DeviceObject,
						 __in PIRP Irp,
						 __in PVOID Context
						 )
{
	PKEVENT Event = (PKEVENT) Context;
	UNREFERENCED_PARAMETER(DeviceObject);
	
	if (Irp->PendingReturned) {

		KeSetEvent(Event, IO_NO_INCREMENT, FALSE);
	}

	return STATUS_MORE_PROCESSING_REQUIRED;
	
}



NTSTATUS
SendSyncIrp(
			__in PDEVICE_OBJECT DeviceObject,
			__in PIRP Irp
			)
{
	KEVENT Event;
	NTSTATUS ns;

	KeInitializeEvent(&Event, NotificationEvent, FALSE);

	IoCopyCurrentIrpStackLocationToNext(Irp);

	IoSetCompletionRoutine(
		Irp, 
		SyncIrpCompletionRoutine,
		&Event, 
		TRUE, 
		TRUE, 
		TRUE);

	ns = IoCallDriver(DeviceObject, Irp);
	if (ns == STATUS_PENDING) {

		KeWaitForSingleObject(
			&Event, 
			Executive, 
			KernelMode, 
			FALSE, 
			NULL);

		ns = Irp->IoStatus.Status;
	}

	return ns;
}


NTSTATUS
DispatchPassThrough(
					__in PDEVICE_OBJECT DeviceObject,
					__in PIRP Irp
					)
{
	PDISK_FLT_DEVICE_EXT DevExt = DeviceObject->DeviceExtension;
	return SkipAndCallDriver(DevExt->LowerDeviceObject, Irp);
}


NTSTATUS
DispatchPower(
			  __in PDEVICE_OBJECT DeviceObject,
			  __in PIRP Irp
			  )
{

	PDISK_FLT_DEVICE_EXT DevExt = DeviceObject->DeviceExtension;
#if (NTDDI_VERSION < NTDDI_VISTA)
	PoStartNextPowerIrp(Irp);
	IoSkipCurrentIrpStackLocation(Irp);
	return PoCallDriver(DevExt->LowerDeviceObject, Irp);
#else
	return SkipAndCallDriver(DevExt->LowerDeviceObject, Irp);
#endif  
}


NTSTATUS
DispatchPnp(
			__in PDEVICE_OBJECT DeviceObject,
			__in PIRP Irp
			)
{
	PDISK_FLT_DEVICE_EXT DevExt = DeviceObject->DeviceExtension;
	PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS ns;

	switch (IrpSp->MinorFunction) {

		case IRP_MN_REMOVE_DEVICE: 
			{
				StartFilter = FALSE;

				if (DevExt->LowerDeviceObject != NULL) {

					IoDetachDevice(DevExt->LowerDeviceObject);
					DevExt->LowerDeviceObject = NULL;
				}

				if (DevExt->ThreadObject != NULL && DevExt->ThreadFlag != TRUE) {

					DevExt->ThreadFlag = TRUE;
					KeSetEvent(&DevExt->ReqEvent,
						(KPRIORITY)0,
						FALSE);

					KeWaitForSingleObject(
						DevExt->ThreadObject,
						Executive,
						KernelMode,
						FALSE,
						NULL
						);

					ObDereferenceObject(DevExt->ThreadObject);

					DevExt->ThreadObject = NULL;
				}

				ExDeletePagedLookasideList(&DevExt->SectorBufHeader);
				ExDeletePagedLookasideList(&DevExt->WriteElemHeader);

				if (DevExt->FilterDeviceObject != NULL) {

					IoDeleteDevice(DevExt->FilterDeviceObject);
				}

				if (GlobalContext.FileHandle != NULL) {

					CloseMapping(GlobalContext.FileHandle,
						GlobalContext.MapBuffer);

					GlobalContext.FileHandle = NULL;
				}

				if (GlobalContext.SparseFileHandle != NULL) {

					ZwClose(GlobalContext.SparseFileHandle);
					GlobalContext.SparseFileHandle = NULL;
				}

			}
			break;

		case IRP_MN_DEVICE_USAGE_NOTIFICATION: 
			{
				if (IrpSp->Parameters.UsageNotification.Type != DeviceUsageTypePaging) {

					ns = SkipAndCallDriver(DevExt->LowerDeviceObject, Irp);
					return ns;
				}

				ns = SendSyncIrp(DevExt->LowerDeviceObject, Irp);
				if (NT_SUCCESS(ns)) {
					
					if (DevExt->LowerDeviceObject->Flags & DO_POWER_PAGABLE) {

						DeviceObject->Flags |= DO_POWER_PAGABLE;
					}
					else {

						DeviceObject->Flags &= ~DO_POWER_PAGABLE;
					}

				}

				IoCompleteRequest(Irp, IO_NO_INCREMENT);
				return ns;

			}
			break;

		default:
			break;
	}

	return SkipAndCallDriver(DevExt->LowerDeviceObject, Irp);
}


NTSTATUS
DispatchReadWrite(
				  __in PDEVICE_OBJECT DeviceObject,
				  __in PIRP Irp
				  )
{
	PDISK_FLT_DEVICE_EXT DevExt = DeviceObject->DeviceExtension;
	PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS ns;
	if (StartFilter && InitFilter) {

		if (IrpSp->MajorFunction == IRP_MJ_READ) {

// 			KdPrint(("Read  Offset:%016I64X Length:%08X\n", 
// 				IrpSp->Parameters.Read.ByteOffset.QuadPart,
// 				IrpSp->Parameters.Read.Length));
		}
		else {

// 			KdPrint(("Write Offset:%016I64X Length:%08X\n",
// 				IrpSp->Parameters.Write.ByteOffset.QuadPart,
// 				IrpSp->Parameters.Write.Length));

			ns = BackupSectorAndSetBit(DevExt->PhysicalDeviceObject, 
				GlobalContext.SparseFileHandle, 
				&IrpSp->Parameters.Write.ByteOffset,
				IrpSp->Parameters.Write.Length, 
				DevExt);

			if (!NT_SUCCESS(ns)) {


				KdPrint(("[FAILED] Write Offset:%016I64X Length:%08X\n",
					IrpSp->Parameters.Write.ByteOffset.QuadPart,
					IrpSp->Parameters.Write.Length));
			}
		}
	}
	

	return SkipAndCallDriver(DevExt->LowerDeviceObject, Irp);
}


NTSTATUS
DispatchDeviceControl(
					  __in PDEVICE_OBJECT DeviceObject,
					  __in PIRP Irp
					  )
{
	PDISK_FLT_DEVICE_EXT DevExt = DeviceObject->DeviceExtension;
	NTSTATUS ns;
	PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
	KEVENT Event;
	PIRP SynIrp;
	UCHAR DBR[512];
	ULONG DBRLength = 512;
	IO_STATUS_BLOCK Iob;
	LARGE_INTEGER ReadOffset = {0};

	switch (IrpSp->Parameters.DeviceIoControl.IoControlCode) {

		case IOCTL_VOLUME_ONLINE:
			{

				KeInitializeEvent(&Event, NotificationEvent, FALSE);

				IoCopyCurrentIrpStackLocationToNext(Irp);
				IoSetCompletionRoutine(Irp, 
					SyncIrpCompletionRoutine, 
					&Event,
					TRUE,
					TRUE,
					TRUE);

				ns = IoCallDriver(DevExt->LowerDeviceObject, Irp);

				if (ns == STATUS_PENDING) {

					KeWaitForSingleObject(&Event,
						Executive,
						KernelMode,
						FALSE,
						NULL);
				}

				KeClearEvent(&Event);
				SynIrp = IoBuildAsynchronousFsdRequest(IRP_MJ_READ,
					DevExt->PhysicalDeviceObject,
					DBR,
					DBRLength,
					&ReadOffset,
					&Iob);

				if (SynIrp != NULL) {
					
					IoSetCompletionRoutine(SynIrp,
						SyncIrpCompletionRoutine,
						&Event,
						TRUE,
						TRUE,
						TRUE);

					ns = IoCallDriver(DevExt->PhysicalDeviceObject, SynIrp);
					if(ns == STATUS_PENDING) {

						KeWaitForSingleObject(&Event,
							Executive,
							KernelMode,
							FALSE,
							NULL);
					}

					ns = SynIrp->IoStatus.Status;

					IoFreeIrp(SynIrp);
					if (NT_SUCCESS(ns)) {

						RtlCopyMemory(&BootSec, DBR, sizeof(NTFS_BOOT_SEC));

						BytesPerSector = BootSec.Bpb.BytesPerSector;

						ExInitializePagedLookasideList(&DevExt->SectorBufHeader, 
							NULL, 
							NULL, 
							0, 
							BytesPerSector, 
							'back', 
							0);

						StartFilter = TRUE;
					}

				}				

				IoCompleteRequest(Irp, IO_NO_INCREMENT);
				return ns;
			}
			break;
		
		default:
			break;
	}
	return SkipAndCallDriver(DevExt->LowerDeviceObject, Irp);
}


NTSTATUS
__fastcall
BackupSectorAndSetBit(
					  __in PDEVICE_OBJECT PhysicalDeviceObject,
					  __in HANDLE File,
					  __in PLARGE_INTEGER Offset,
					  __in ULONG Length,
					  __in PDISK_FLT_DEVICE_EXT DevExt
					  )
{
	ULONGLONG BitOffset = BYTEOFF_TO_BITOFF(Offset->QuadPart);
	ULONG BitLength = BYTELEN_TO_BITLEN(Length);
	ULONG i;
	ULONG BitStatus;
	BOOLEAN BitRet;
	PUCHAR BitBuf = (PUCHAR)GlobalContext.MapBuffer;
	ULONG BitSize = GlobalContext.MapSize;
	LARGE_INTEGER BackupOffset;
	PUCHAR BackupBuf;
	NTSTATUS ns;
	PWRITE_ELEMENT WriteElem;

	ASSERT((Offset->QuadPart % BytesPerSector) == 0);
	ASSERT((Length % BytesPerSector) == 0);

	for (i = 0; i < BitLength; i++) {

		BitStatus = GetBit(BitBuf, BitSize, BitOffset + i);
		if (BitStatus == -1) {

			return STATUS_UNSUCCESSFUL;
		}
		else if (BitStatus == 0) {

			BackupBuf = ExAllocateFromPagedLookasideList(&DevExt->SectorBufHeader);

			BackupOffset.QuadPart = (i + BitOffset) * BytesPerSector;
			ns = DirectReadWrite(PhysicalDeviceObject, 
				DIRECT_READ, 
				BackupBuf, 
				&BackupOffset, 
				BytesPerSector);
			if (!NT_SUCCESS(ns)) {

				ExFreeToPagedLookasideList(&DevExt->SectorBufHeader, BackupBuf);
				return ns;
			}

			WriteElem = (PWRITE_ELEMENT)ExAllocateFromPagedLookasideList(&DevExt->WriteElemHeader);
			if (WriteElem == NULL) {

				ExFreeToPagedLookasideList(&DevExt->SectorBufHeader, BackupBuf);
				return STATUS_INSUFFICIENT_RESOURCES;
			}

			WriteElem->File = File;
			WriteElem->Buffer = BackupBuf;
			WriteElem->Length = BytesPerSector;
			WriteElem->Offset.QuadPart = BackupOffset.QuadPart;
			WriteElem->NeedFree = TRUE;
			WriteElem->Tag = 'back';

			ExInterlockedInsertTailList(&DevExt->ReqList,
				&WriteElem->ListEntry,
				&DevExt->ReqLock);

			KeSetEvent(&DevExt->ReqEvent, 
				(KPRIORITY)0, 
				FALSE);

			BitRet = SetBit(BitBuf, BitSize, BitOffset + i);
			if (!BitRet) {

				return STATUS_UNSUCCESSFUL;
			}
		}
	}

	FlushBuffer(DevExt,
		GlobalContext.FileHandle, 
		BitBuf, 
		BITOFF_TO_MAPOFF(BitOffset), 
		BITLEN_TO_MAPLEN(BitOffset, BitLength));

	return STATUS_SUCCESS;
}


NTSTATUS
FlushBuffer(
			__in PDISK_FLT_DEVICE_EXT DevExt,
			__in HANDLE File,
			__in PVOID Buffer,
			__in ULONG Offset,
			__in ULONG Size
			)
{
	PUCHAR FlushBuf = (PUCHAR)Buffer;

	PWRITE_ELEMENT WriteElem = (PWRITE_ELEMENT)ExAllocateFromPagedLookasideList(&DevExt->WriteElemHeader);
	if (WriteElem == NULL) {

		return STATUS_INSUFFICIENT_RESOURCES;
	}


	WriteElem->File = File;
	WriteElem->Buffer = FlushBuf + Offset;
	WriteElem->Length = Size;
	WriteElem->Offset.QuadPart = (LONGLONG)Offset;
	WriteElem->NeedFree = FALSE;

	ExInterlockedInsertTailList(&DevExt->ReqList,
		&WriteElem->ListEntry,
		&DevExt->ReqLock);

	KeSetEvent(&DevExt->ReqEvent, 
		(KPRIORITY)0, 
		FALSE);

	return STATUS_SUCCESS;
}



NTSTATUS
DirectReadWrite(
				__in PDEVICE_OBJECT DeviceObject,
				__in ULONG ReadWrite,
				__out PVOID Buffer,
				__in PLARGE_INTEGER Offset,
				__in ULONG Length
				)
{
	NTSTATUS ns;
	ULONG MjFunc;
	KEVENT Event;
	PIRP Irp;
	IO_STATUS_BLOCK Iob;
	PMDL NextMdl;

	if (ReadWrite == DIRECT_READ) {

		MjFunc = IRP_MJ_READ;
	}
	else if (ReadWrite == DIRECT_WRITE) {

		MjFunc = IRP_MJ_WRITE;
	}
	else {

		return STATUS_INVALID_PARAMETER;
	}

	KeInitializeEvent(&Event, NotificationEvent, FALSE);

	Irp = IoBuildAsynchronousFsdRequest(MjFunc,
		DeviceObject,
		Buffer,
		Length,
		Offset,
		&Iob);

	if (Irp == NULL) {

		return STATUS_INSUFFICIENT_RESOURCES;
	}

	IoSetCompletionRoutine(Irp,
		SyncIrpCompletionRoutine,
		&Event,
		TRUE,
		TRUE,
		TRUE);

	ns = IoCallDriver(DeviceObject, Irp);
	if(ns == STATUS_PENDING) {

		KeWaitForSingleObject(&Event,
			Executive,
			KernelMode,
			FALSE,
			NULL);
	}

	ns = Irp->IoStatus.Status;

	while (Irp->MdlAddress != NULL) {

		NextMdl = Irp->MdlAddress->Next;
		MmUnlockPages(Irp->MdlAddress);
		IoFreeMdl(Irp->MdlAddress);

		Irp->MdlAddress = NextMdl;
	}

	IoFreeIrp(Irp);

	return ns;
}

LONGLONG GetFileSize(
					 __in HANDLE Handle
					 )
{
	IO_STATUS_BLOCK Iob;
	FILE_STANDARD_INFORMATION FileStdInfo;
	NTSTATUS ns;

	ns = ZwQueryInformationFile(Handle, 
		&Iob, 
		&FileStdInfo, 
		sizeof(FILE_STANDARD_INFORMATION), 
		FileStandardInformation);

	if (!NT_SUCCESS(ns)) {

		return 0;
	}

	return FileStdInfo.EndOfFile.QuadPart;
}

BOOLEAN
__fastcall
SetBit(
	   __in PUCHAR BitBuf, 
	   __in ULONG BitSize, 
	   __in ULONGLONG Bit
	   )
{
	ULONG ByteOffset = 0;
	ULONG BitOffset = 0;

	ByteOffset = (ULONG)(Bit / 8);
	if (ByteOffset >= BitSize) {

		return FALSE;
	}

	BitOffset = (ULONG)(Bit % 8);

	BitBuf[ByteOffset] |= (UCHAR)(1 << BitOffset);

	return TRUE;
}

BOOLEAN 
__fastcall
ClearBit(
		 __in PUCHAR BitBuf, 
		 __in ULONG BitSize, 
		 __in ULONGLONG Bit 
		 )
{
	ULONG ByteOffset = 0;
	ULONG BitOffset = 0;

	ByteOffset = (ULONG)(Bit / 8);
	if (ByteOffset >= BitSize) {

		return FALSE;
	}

	BitOffset = (ULONG)(Bit % 8);

	BitBuf[ByteOffset] &= (UCHAR)(~(1 << BitOffset));

	return TRUE;
}

ULONG 
__fastcall
GetBit(
	   __in PUCHAR BitBuf, 
	   __in ULONG BitSize,
	   __in ULONGLONG Bit
	   )
{
	ULONG ByteOffset = 0;
	ULONG BitOffset = 0;

	ByteOffset = (ULONG)(Bit / 8);
	if (ByteOffset >= BitSize) {

		return (ULONG)-1;
	}

	BitOffset = (ULONG)(Bit % 8);

	return (BitBuf[ByteOffset] & (UCHAR)(1 << BitOffset)) != 0;
} 


NTSTATUS
OpenMapping(
			__in PWCHAR Name,
			__out PHANDLE File,
			__out PVOID *Buffer,
			__out ULONG *BufferSize
			)
{
	NTSTATUS ns;
	HANDLE FileHandle;
	UNICODE_STRING FileName;
	OBJECT_ATTRIBUTES oa;
	IO_STATUS_BLOCK Iob;
	LARGE_INTEGER MaxSize;
	PVOID MapBuffer;
	LARGE_INTEGER Offset = {0};

	RtlInitUnicodeString(&FileName, Name);
	InitializeObjectAttributes(&oa, &FileName, OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE, NULL, NULL);

	ns = ZwCreateFile(&FileHandle, 
		FILE_READ_DATA | FILE_WRITE_DATA | SYNCHRONIZE, 
		&oa, 
		&Iob,
		NULL,
		0,
		FILE_SHARE_READ,
		FILE_OPEN, 
		FILE_SYNCHRONOUS_IO_NONALERT | FILE_WRITE_THROUGH,
		NULL, 
		0);

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	MaxSize.QuadPart = GetFileSize(FileHandle);
	if (MaxSize.QuadPart == 0) {

		ZwClose(FileHandle);
		return STATUS_UNSUCCESSFUL;
	}

	MapBuffer = ExAllocatePoolWithTag(PagedPool, MaxSize.LowPart, 'bmp ');
	if (MapBuffer == NULL) {

		ZwClose(FileHandle);
		return STATUS_INSUFFICIENT_RESOURCES;
	}

	ns = ZwReadFile(FileHandle, NULL, NULL, NULL, &Iob, MapBuffer, MaxSize.LowPart, &Offset, NULL);
	if (!NT_SUCCESS(ns)) {

		ExFreePoolWithTag(MapBuffer, 'bmp ');
		ZwClose(FileHandle);
		return ns;
	}


	*File = FileHandle;
	*Buffer = MapBuffer;
	*BufferSize = MaxSize.LowPart;

	return ns;
}

VOID
CloseMapping(
			 __in HANDLE File,
			 __in PVOID Buffer
			 )
{
	ExFreePoolWithTag(Buffer, 'bmp ');
	ZwClose(File);
}


NTSTATUS
CreateSparseFile(
				 __out PHANDLE Handle, 
				 __in PWCHAR FileName,
				 __in ULONGLONG FileSize,
				 __in BOOLEAN Restore
				 )
{
	NTSTATUS ns;
	HANDLE FileHandle;
	UNICODE_STRING SparseFileName;
	IO_STATUS_BLOCK	Iob;
	OBJECT_ATTRIBUTES oa;
	FILE_END_OF_FILE_INFORMATION FileEndInfo = {0};
	ULONG ExtFlag = 0;

	if (!Restore) {

		ExtFlag = FILE_RANDOM_ACCESS | FILE_NO_INTERMEDIATE_BUFFERING | FILE_WRITE_THROUGH;
	}

	RtlInitUnicodeString(&SparseFileName, FileName);
	InitializeObjectAttributes(&oa, 
		&SparseFileName,
		OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
		NULL,
		NULL);

	ns = ZwCreateFile(
		&FileHandle,
		FILE_READ_DATA | FILE_WRITE_DATA | SYNCHRONIZE | DELETE,
		&oa,
		&Iob,
		NULL,
		FILE_ATTRIBUTE_NORMAL,
		0,
		FILE_OPEN_IF,
		FILE_SYNCHRONOUS_IO_NONALERT | ExtFlag,
		NULL,
		0);

	if(!NT_SUCCESS(ns)) {

		return ns;
	}

	if (Iob.Information == FILE_CREATED) {

		ns = ZwFsControlFile(FileHandle,
			NULL,
			NULL,
			NULL,
			&Iob,
			FSCTL_SET_SPARSE,
			NULL,
			0,
			NULL,
			0);

		if(!NT_SUCCESS(ns)) {

			ZwClose(FileHandle);
			return ns;
		}

		FileEndInfo.EndOfFile.QuadPart = FileSize + 10*1024*1024;
		ns = ZwSetInformationFile(FileHandle,
			&Iob,
			&FileEndInfo,
			sizeof(FILE_END_OF_FILE_INFORMATION),
			FileEndOfFileInformation);

		if (!NT_SUCCESS(ns)) {
			ZwClose(FileHandle);
			return ns;
		}
	}
	
	*Handle = FileHandle;
	return ns;
}


NTSTATUS
AddDevice(
		  __in PDRIVER_OBJECT	DriverObject,
		  __in PDEVICE_OBJECT	PhysicalDeviceObject
		  )
{
	NTSTATUS ns;
	PDISK_FLT_DEVICE_EXT DevExt;
	PDEVICE_OBJECT LowerDevObj;
	PDEVICE_OBJECT FltDevObj;
	ULONG RetLength;
	HANDLE ThreadHandle;
	UCHAR NameBuffer[1024 + sizeof(OBJECT_NAME_INFORMATION)];
	POBJECT_NAME_INFORMATION DeviceNameInfo = (POBJECT_NAME_INFORMATION)NameBuffer;

	DECLARE_CONST_UNICODE_STRING(Vol1, L"\\Device\\HarddiskVolume1");

	ns = ObQueryNameString(PhysicalDeviceObject, DeviceNameInfo, 1024, &RetLength);
	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	if (RtlCompareUnicodeString(&Vol1, &DeviceNameInfo->Name, TRUE) != 0) {

		return STATUS_SUCCESS;
	}

	KdPrint(("Physical Device Dos Name : %wZ\n", &DeviceNameInfo->Name));

	ns = IoCreateDevice(DriverObject,
		sizeof(DISK_FLT_DEVICE_EXT),
		NULL,
		FILE_DEVICE_DISK,
		FILE_DEVICE_SECURE_OPEN,
		FALSE,
		&FltDevObj);

	if (!NT_SUCCESS(ns)) {

		return ns;
	}

	DevExt = FltDevObj->DeviceExtension;
	RtlZeroMemory(DevExt,sizeof(DISK_FLT_DEVICE_EXT));

	LowerDevObj = IoAttachDeviceToDeviceStack(FltDevObj, 
		PhysicalDeviceObject);

	if (LowerDevObj == NULL) {

		IoDeleteDevice(FltDevObj);
		ns = STATUS_NO_SUCH_DEVICE;
		return ns;
	}


	FltDevObj->Flags |= LowerDevObj->Flags & (DO_DIRECT_IO | DO_BUFFERED_IO | DO_POWER_PAGABLE);
	FltDevObj->Flags &= ~DO_DEVICE_INITIALIZING;


	DevExt->DeviceType = DEVICE_TYPE_FILTER;
	DevExt->FilterDeviceObject = FltDevObj;
	DevExt->PhysicalDeviceObject = PhysicalDeviceObject;
	DevExt->LowerDeviceObject = LowerDevObj;



	InitializeListHead(&DevExt->ReqList);
	KeInitializeSpinLock(&DevExt->ReqLock);

	KeInitializeEvent(&DevExt->ReqEvent,
		SynchronizationEvent,
		FALSE);
	DevExt->ThreadFlag = FALSE;

	ns = PsCreateSystemThread(&ThreadHandle,
		(ACCESS_MASK)0L,
		NULL,
		NULL,
		NULL,
		WriteWorkThread,
		DevExt);

	if (!NT_SUCCESS(ns)) {

		IoDeleteDevice(FltDevObj);
		return ns;
	}

	ns = ObReferenceObjectByHandle(ThreadHandle,
		THREAD_ALL_ACCESS,
		NULL,
		KernelMode,
		&DevExt->ThreadObject,
		NULL);

	ZwClose(ThreadHandle);

	if (!NT_SUCCESS(ns)) {
		DevExt->ThreadFlag = TRUE;
		KeSetEvent(&DevExt->ReqEvent,
			(KPRIORITY)0,
			FALSE);
		
		IoDeleteDevice(FltDevObj);
		return ns;
	}

	ExInitializePagedLookasideList(&DevExt->WriteElemHeader, 
		NULL, 
		NULL, 
		0, 
		sizeof(WRITE_ELEMENT), 
		'welm', 
		0);


	return ns;
}

VOID
Unload(
	   __in PDRIVER_OBJECT	DriverObject
	   )
{
	PDEVICE_OBJECT NextDevObj;
	PDEVICE_OBJECT CurDevObj;
	PDISK_FLT_DEVICE_EXT DevExt;

	NextDevObj = DriverObject->DeviceObject;

	while (NextDevObj != NULL) {


		CurDevObj = NextDevObj;
		DevExt = (PDISK_FLT_DEVICE_EXT)CurDevObj->DeviceExtension;
		NextDevObj = CurDevObj->NextDevice;

		if (DevExt->DeviceType == DEVICE_TYPE_FILTER) {

			IoDetachDevice(DevExt->LowerDeviceObject);

			if (DevExt->ThreadObject != NULL && DevExt->ThreadFlag != TRUE) {

				DevExt->ThreadFlag = TRUE;
				KeSetEvent(&DevExt->ReqEvent,
					(KPRIORITY)0,
					FALSE);

				KeWaitForSingleObject(
					DevExt->ThreadObject,
					Executive,
					KernelMode,
					FALSE,
					NULL
					);

				ObDereferenceObject(DevExt->ThreadObject);

				DevExt->ThreadObject = NULL;
			}

			ExDeletePagedLookasideList(&DevExt->SectorBufHeader);
			ExDeletePagedLookasideList(&DevExt->WriteElemHeader);
			DevExt->LowerDeviceObject = NULL;
		}
		
		IoDeleteDevice(CurDevObj);
	}

	if (GlobalContext.FileHandle != NULL) {

		CloseMapping(GlobalContext.FileHandle,
			GlobalContext.MapBuffer);

		GlobalContext.FileHandle = NULL;
	}

	if (GlobalContext.SparseFileHandle != NULL) {

		ZwClose(GlobalContext.SparseFileHandle);
		GlobalContext.SparseFileHandle = NULL;
	}


	return;
}


VOID
ReinitializationRoutine(
						__in PDRIVER_OBJECT DriverObject, 
						__in PVOID Context, 
						__in ULONG Count
						)
{
	NTSTATUS ns;
	PDEVICE_OBJECT FilterDevice;
	PDISK_FLT_DEVICE_EXT DevExt;
	BOOLEAN PrepareRestore = IsRestoreDisk();

	UNREFERENCED_PARAMETER(DriverObject);
	UNREFERENCED_PARAMETER(Context);
	UNREFERENCED_PARAMETER(Count);

	ns = OpenMapping(L"\\??\\D:\\Bitmap.dat", 
		&GlobalContext.FileHandle, 
		&GlobalContext.MapBuffer,
		&GlobalContext.MapSize);

	if (!NT_SUCCESS(ns)) {

		return;
	}

	ns = CreateSparseFile(&GlobalContext.SparseFileHandle, 
		L"\\??\\D:\\SpareFile.dat", 
		BootSec.ExBpb.TotalSectors * BootSec.Bpb.BytesPerSector, PrepareRestore);

	if (!NT_SUCCESS(ns)) {

		return;
	}

	if (PrepareRestore) {

		FilterDevice = DriverObject->DeviceObject;

		while (FilterDevice != NULL) {

			DevExt = FilterDevice->DeviceExtension;
			if (DevExt != NULL && DevExt->DeviceType == DEVICE_TYPE_FILTER) {

				RestoreDisk(DevExt);
				KdBreakPoint();
			}

			FilterDevice = FilterDevice->NextDevice;
		}
	}

	InitFilter = TRUE;

	return;
}

NTSTATUS 
DriverEntry(
			__in PDRIVER_OBJECT DriverObject, 
			__in PUNICODE_STRING RegistryPath 
			)
{
	ULONG i;

	UNREFERENCED_PARAMETER(RegistryPath);
	KdBreakPoint();

	for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++) {

		DriverObject->MajorFunction[i] = DispatchPassThrough;
	}


	DriverObject->MajorFunction[IRP_MJ_POWER] = DispatchPower;
	DriverObject->MajorFunction[IRP_MJ_PNP] = DispatchPnp;
	DriverObject->MajorFunction[IRP_MJ_READ] = DispatchReadWrite;
	DriverObject->MajorFunction[IRP_MJ_WRITE] = DispatchReadWrite;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DispatchDeviceControl;

	DriverObject->DriverExtension->AddDevice = AddDevice;
	DriverObject->DriverUnload = Unload;

	IoRegisterBootDriverReinitialization(DriverObject,
		ReinitializationRoutine,
		NULL);

	return STATUS_SUCCESS;
}